import triton
import triton.language as tl
import test_common
import pytest
import torch
import torch_npu
import numpy as np
import random

@triton.jit
def softmax_triton(
    x_ptr,
    x_row_stride,
    y_ptr,
    y_row_stride,
    n_rows,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row_start = tl.program_id(0)
    row_step = tl.num_programs(0)
    for row_idx in tl.range(row_start, n_rows, row_step):
        x_ptr_base = x_ptr + row_idx * x_row_stride
        y_ptr_base = y_ptr + row_idx * y_row_stride
        col_offsets = tl.arange(0, BLOCK_SIZE)
        input_ptrs = x_ptr_base + col_offsets
        mask = col_offsets < n_cols
        row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
        row = row.to(tl.float32)
        row_minus_max = row - tl.max(row, axis=0)
        numerator = tl.exp(row_minus_max)
        denominator = tl.sum(numerator, axis=0)
        softmax_output = numerator / denominator
        output_ptrs = y_ptr_base + col_offsets
        tl.store(output_ptrs, softmax_output.to(y_ptr.dtype.element_ty), mask=mask)

def softmax(x):
    n_rows, n_cols = x.shape
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x).npu()
    num_programs = 40
    num_programs = min(num_programs, n_rows)
    softmax_triton[(num_programs, 1, 1)](
        y,
        x,
        x.stride(0),
        y.stride(0),
        n_rows,
        n_cols,
        BLOCK_SIZE
    )
    return y